package tech.waterism.algorithm.moea.PSO;

import cn.hutool.extra.spring.SpringUtil;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.BeanUtils;
import tech.waterism.constant.DispatchConstants;
import tech.waterism.constant.StringConstants;
import tech.waterism.document.DispatchMoeaDocument;
import tech.waterism.dto.*;
import tech.waterism.enums.FieldEnum;
import tech.waterism.enums.LogicEnum;
import tech.waterism.service.DispatchService;
import tech.waterism.utils.moea.DispatchUtils;

import java.util.*;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicReference;
import java.util.stream.Collectors;

@Data
@Slf4j
public class DispatchFitnessCO {

    private static final double AEQ_LIMIT = 0.01;
    private final TopoDataDto topoDataDto;
    private final AlgorithmDispatch algorithm;
    private final double[] paras;
    private final DispatchMoeaDocument document;
    private final DispatchService dispatchService;
    private final Map<String, DispatchConstraint> QConstraintsMap;
    private List<String> message;

    public DispatchFitnessCO(TopoDataDto topoDataDto, AlgorithmDispatch algorithm, double[] paras,
                             DispatchMoeaDocument document) {
        this.topoDataDto = topoDataDto;
        this.algorithm = algorithm;
        this.paras = paras;
        this.message = new LinkedList<>();
        this.document = document;
        dispatchService = SpringUtil.getBean(DispatchService.class);
        Map<String, TopoObjectDto> topoObjectDtoMap = topoDataDto.getTopoObject().stream()
                .collect(Collectors.toMap(i -> i.getFpIndex().getStcd(), v -> v, (k1, k2) -> k1));
        for (DispatchConstraint constrain : algorithm.getConstrains()) {
            String stcd = constrain.getStcd();
            if (topoObjectDtoMap.containsKey(stcd)) {
                constrain.setStnm(topoObjectDtoMap.get(stcd).getFpIndex().getStnm());
            }
        }
        QConstraintsMap = algorithm.getConstrains().stream()
                .filter(i -> (i.getLogic().equals(LogicEnum.LT) || i.getLogic().equals(LogicEnum.LTE)) &&
                        i.getCompareField().equals(FieldEnum.QSIM))
                .collect(Collectors.toMap(DispatchConstraint::getStcd, i -> i));
    }

    public double[] getFitness() throws Exception {
        int oriSize = algorithm.getObjectives().size();
        AtomicInteger numberOfObjectivesAppend = new AtomicInteger();
        this.document.getExcessMap().forEach((k, v) -> {
            numberOfObjectivesAppend.addAndGet(v.getWaterMap().size());
        });

        double[] result = new double[oriSize + numberOfObjectivesAppend.get()];
        try {
            topoDataDto.setDispatchType(DispatchConstants.DISPATCH_PLAN_TYPE_BEGINOUTFLOW);
            // 全开之后，计算原始目标值
            List<TopoDataDto> firstResults = dispatchService.calculate(topoDataDto);
            boolean checkResult = CheckUtil.checkRz(firstResults, algorithm.getConstrains());
            if (checkResult) {
                firstResults = dispatchService.calculate(firstResults.get(0));
            }
            Map<String, TopoObjectDto> topoObjectDtoMap = firstResults.get(0).getTopoObject().stream()
                    .collect(Collectors.toMap(i -> i.getFpIndex().getStcd(), v -> v, (k1, k2) -> k1));
            for (int i = 0; i < algorithm.getObjectives().size(); i++) {
                DispatchObjective objective = algorithm.getObjectives().get(i);
                if (topoObjectDtoMap.containsKey(objective.getStcd())) {
                    result[i] = Optional.ofNullable(
                                    CheckUtil.calFitness(topoObjectDtoMap.get(objective.getStcd()).getFpIndex(),
                                            topoDataDto.getFymdh(), objective))
                            .orElseThrow(() -> new Exception("计算目标值失败"));
                }
                if (objective.getStcd().contains(StringConstants.COMMA)) {
                    String[] stcds = objective.getStcd().split(StringConstants.COMMA);
                    if (objective.getCompareField().equals(FieldEnum.POWER_E)) {
                        List<FPIndexDto> indexDtos = new ArrayList<>();
                        for (String stcd : stcds) {
                            if (topoObjectDtoMap.containsKey(stcd)) {
                                indexDtos.add(topoObjectDtoMap.get(stcd).getFpIndex());
                            }
                        }
                        result[i] = Optional.ofNullable(
                                        CheckUtil.calFitnessCombined(indexDtos, topoDataDto.getFymdh(), objective))
                                .orElseThrow(() -> new Exception("计算目标值失败"));
                    }
                }
            }
            // 目标追加的index，从原有目标的最后一个开始
            int objectiveIndex = algorithm.getObjectives().size();
            // 水库组顺次全关闭，追加目标值
            for (Map.Entry<Integer, ExcessQDto> entry : this.document.getExcessMap().entrySet()) {
                Integer index = entry.getKey();

                List<DateInterval> intervalList = entry.getValue().getIntervalList();
                List<List<String>> groupList = new LinkedList<>();
                for (WaterAlphaDto waterAlphaDto : entry.getValue().getAlphaList()) {
                    groupList.add(waterAlphaDto.getStcds());
                }
                for (List<String> stringList : groupList) {
                    AtomicReference<Double> sumSingleWater = new AtomicReference<>(0.0);
                    AtomicReference<Double> sumQa = new AtomicReference<>(0.0);
                    TopoDataDto calDto1 = new TopoDataDto();
                    BeanUtils.copyProperties(topoDataDto, calDto1);
                    calDto1.setDispatchType(DispatchConstants.DISPATCH_PLAN_TYPE_OPTIMIZEDOUTFLOW);
                    generCalData(calDto1, stringList);
                    List<TopoDataDto> tempResults1 = dispatchService.calculate(calDto1);
                    boolean checkResult1 = CheckUtil.checkRz(tempResults1, algorithm.getConstrains());
                    if (checkResult1) {
                        tempResults1 = dispatchService.calculate(tempResults1.get(0));
                    }
                    TopoDataDto tempResult1 = tempResults1.get(0);
                    TopoObjectDto controlNode1 = tempResult1.getTopoObject().get(index);
                    DispatchConstraint constraint = QConstraintsMap.get(controlNode1.getFpIndex().getStcd());
                    LogFieldDto logFieldDto = new LogFieldDto().setStcd(controlNode1.getFpIndex().getStcd())
                            .setStnm(controlNode1.getFpIndex().getStnm());
                    for (ForecastDataDto item : controlNode1.getFpIndex().getData()) {
                        intervalList.forEach(interval -> {
                            if (!item.getDt().before(interval.getStart()) && !item.getDt().after(interval.getEnd())) {
                                sumSingleWater.updateAndGet(v -> v + item.getQSim());
                                sumQa.updateAndGet(v -> {
                                    try {
                                        return v + getQa(logFieldDto, item, constraint,
                                                controlNode1.getFpIndex().getFloodCharacter());
                                    } catch (Exception e) {
                                        e.printStackTrace();
                                        return (double) 0;
                                    }
                                });
                            }
                        });
                    }
                    double sumTemp = sumSingleWater.get();
                    result[objectiveIndex++] = Math.pow((sumTemp - sumQa.get()), 2);
                }
            }

            PunishDto punishDto = punish();
            this.message = punishDto.getMessage();
            double punish = punishDto.getPunish();
            for (int i = 0; i < result.length; i++) {
                result[i] = result[i] + result[i] * 1e4 * punish;
            }
            log.debug("punish:" + punish);
            // 归一化
            for (int i = 0; i < algorithm.getObjectives().size(); i++) {
                result[i] = result[i] * algorithm.getObjectives().get(i).getWeight();
            }
            // 追加的目标值不归一化，直接按照惩罚值计算
            return result;
        } catch (Exception e) {
            log.error(e.getMessage(), e);
            this.message = Arrays.asList(e.getMessage());
            return null;
        }
    }

    private Double getQa(LogFieldDto logFieldDto, ForecastDataDto datum, DispatchConstraint constrain,
                         FloodCharacterDto floodCharacter) throws Exception {
        double diff = 0;
        String compareField = constrain.getCompareField().getField();
        Object value;
        switch (constrain.getValueType()) {
            case FIXED:
                value = CheckUtil.obj2Map(datum).get(compareField);
                diff = CheckUtil.getConstraintDiff(constrain.getValue(), (Double) value, constrain.getLogic());
                break;
            case CHARACTER:
                Double constrainValue =
                        CheckUtil.getCharacterValue(logFieldDto, constrain.getCatalog(), floodCharacter, datum);
                value = CheckUtil.obj2Map(datum).get(compareField);
                diff = CheckUtil.getConstraintDiff(constrainValue, (Double) value, constrain.getLogic());
                break;
            default:
                throw new Exception("约束值类型错误");
        }
        return diff;
    }

    private void generCalData(TopoDataDto calDto1, List<String> stringList) {
        for (TopoObjectDto topoObjectDto : calDto1.getTopoObject()) {
            if (stringList.contains(topoObjectDto.getFpIndex().getStcd())) {
                DispatchUtils.setQoutZeroSingle(topoObjectDto, topoDataDto.getFymdh());
            }
        }
    }

    /**
     * get punish value by constraints
     */
    private PunishDto punish() throws Exception {
        List<DispatchConstraint> constrains = algorithm.getConstrains();
        if (constrains == null || constrains.size() == 0) {
            return null;
        }
        PunishDto result = new PunishDto();
        List<String> message = new LinkedList<>();
        Date fymdh = topoDataDto.getFymdh();
        Map<String, TopoObjectDto> topoObjectDtoMap = topoDataDto.getTopoObject().stream()
                .collect(Collectors.toMap(i -> i.getFpIndex().getStcd(), v -> v, (k1, k2) -> k1));
        double checkResult = 0;
        for (DispatchConstraint constrain : constrains) {
            String stcd = constrain.getStcd();
            String stnm = constrain.getStnm();
            double checkResultItem = 0;
            if (topoObjectDtoMap.containsKey(stcd)) {
                FPIndexDto fpIndexDto = topoObjectDtoMap.get(stcd).getFpIndex();
                LogFieldDto logFieldDto = new LogFieldDto().setStcd(stcd).setStnm(stnm);
                List<ForecastDataDto> data = fpIndexDto.getData();
                List<ForecastDataDto> checkData =
                        data.stream().filter(datum -> datum.getDt().after(fymdh)).collect(Collectors.toList());
                if (constrain.getCompareField().equals(FieldEnum.Z_END)) {
                    checkResultItem =
                            CommonFitness.calFieldZend(logFieldDto, checkData.get(checkData.size() - 1), constrain,
                                    fpIndexDto.getFloodCharacter()) * 1e6;
                } else if (constrain.getCompareField().equals(FieldEnum.QOUT_DIFF)) {
                    checkResultItem = calFieldDiff(data, constrain) * 1e2;
                    // } else if (constrain.getCompareField().equals(FieldEnum.RZ)) {
                    //     checkResultItem =
                    //             CommonFitness.calField(logFieldDto, checkData, constrain, fpIndexDto
                    //             .getFloodCharacter()) *
                    //                     1e4;
                } else if (constrain.getCompareField().equals(FieldEnum.POWER_N)) {
                    // checkResultItem = CommonFitness.calFieldPowerN(data, constrain) * 1e6;
                    CommonFitness.calFieldPowerNV2(fpIndexDto, data, constrain);
                } else if (constrain.getCompareField().equals(FieldEnum.POWER_Q)) {
                    checkResultItem = CommonFitness.calFieldPowerQ(data, constrain) * 1e2;
                } else {
                    checkResultItem =
                            CommonFitness.calField(logFieldDto, checkData, constrain, fpIndexDto.getFloodCharacter());
                }
            }
            if (checkResultItem > 0) {
                message.add(stnm + "(" + stcd + ")" + constrain.getCompareField().getName() +
                        constrain.getLogic().getName() + "未满足约束");
                checkResult += checkResultItem;
            }
        }
        result.setPunish(checkResult);
        result.setMessage(message);
        return result;
    }

    private double calFieldDiff(List<ForecastDataDto> checkData, DispatchConstraint constrain) {
        double punish = 0;
        int segIndex = constrain.getSegInfo().getSegIndex();
        int varIndex = constrain.getSegInfo().getVarIndex();
        System.out.println(checkData.stream().map(ForecastDataDto::getQOut).collect(Collectors.toList()));
        if (paras != null) {
            for (int i = 0; i < varIndex - segIndex; i++) {
                if (paras[segIndex + i] >= checkData.size()) {
                    continue;
                }
                double qoutEnd = checkData.get((int) paras[segIndex + i]).getQOut();
                double qoutStart = checkData.get((int) paras[segIndex + i] - 1).getQOut();
                System.out.println("end:" + qoutEnd);
                System.out.println("start:" + qoutStart);
                punish += CheckUtil.getConstraintDiff(constrain.getValue(), Math.abs(qoutEnd - qoutStart),
                        constrain.getLogic());
            }
            return punish;
        } else {
            return 0;
        }
    }

    private boolean checkUsedgatesOverNgates() {
        List<DispatchVariable> gateVariable =
                algorithm.getVariables().stream().filter(v -> v.getCompareField().equals(FieldEnum.GATE))
                        .collect(Collectors.toList());
        if (gateVariable.size() == 0) {
            return false;
        }
        Date fymdh = topoDataDto.getFymdh();
        Map<String, DispatchVariable> variableMap =
                gateVariable.stream().collect(Collectors.toMap(DispatchVariable::getStcd, v -> v, (k1, k2) -> k1));
        for (TopoObjectDto topoObjectDto : topoDataDto.getTopoObject()) {
            String stcd = topoObjectDto.getFpIndex().getStcd();
            if (variableMap.containsKey(stcd)) {
                // DispatchVariable variable = variableMap.get(stcd);
                boolean flg = CommonFitness.checkNGates(topoObjectDto.getFpIndex(), fymdh);
                if (flg) {
                    return true;
                }
            }
        }
        return false;
    }

    public CheckConstrainDto checkConstrain(Boolean pre) throws Exception {
        CheckConstrainDto result = new CheckConstrainDto();
        List<DispatchConstraint> constrains = algorithm.getConstrains();
        if (constrains == null || constrains.size() == 0) {
            result.setFlg(true);
            return result;
        }
        Date fymdh = topoDataDto.getFymdh();
        Map<String, List<DispatchConstraint>> constrainMap =
                constrains.stream().collect(Collectors.groupingBy(DispatchConstraint::getStcd));
        for (TopoObjectDto topoObjectDto : topoDataDto.getTopoObject()) {
            String stcd = topoObjectDto.getFpIndex().getStcd();
            if (constrainMap.containsKey(stcd)) {
                List<DispatchConstraint> constraintList = constrainMap.get(stcd);
                LogFieldDto logFieldDto = new LogFieldDto().setStcd(stcd).setStnm(topoObjectDto.getFpIndex().getStnm());
                List<ForecastDataDto> data = topoObjectDto.getFpIndex().getData();
                List<ForecastDataDto> checkData =
                        data.stream().filter(datum -> datum.getDt().after(fymdh)).collect(Collectors.toList());
                for (DispatchConstraint constrain : constraintList) {
                    if (Boolean.TRUE.equals(pre)) {
                        // 可执行判断，只检查河道流量上限、水位上限和水库入流上限约束
                        if (constrain.getCompareField().equals(FieldEnum.QSIM) ||
                                constrain.getCompareField().equals(FieldEnum.ZSIM) ||
                                constrain.getCompareField().equals(FieldEnum.QIN)) {
                            if (constrain.getLogic() == LogicEnum.LTE || constrain.getLogic() == LogicEnum.LT) {
                                boolean checkResult = CommonFitness.checkField(logFieldDto, checkData, constrain,
                                        topoObjectDto.getFpIndex().getFloodCharacter());
                                if (!checkResult) {
                                    result.setFlg(false);
                                    String message = constrain.getStnm() + "(" + constrain.getStcd() + ")" +
                                            constrain.getCompareField().getName() + constrain.getLogic().getName() +
                                            "未满足约束";
                                    result.setMessage(message);
                                    return result;
                                }
                            }
                        }
                    } else {
                        if (constrain.getCompareField().equals(FieldEnum.Z_END)) {
                            double punish = CommonFitness.calFieldZend(logFieldDto, checkData.get(checkData.size() - 1),
                                    constrain, topoObjectDto.getFpIndex().getFloodCharacter());
                            if (punish > AEQ_LIMIT) {
                                result.setFlg(false);
                                String message = constrain.getStnm() + "(" + constrain.getStcd() + ")" +
                                        constrain.getCompareField().getName() + constrain.getLogic().getName() +
                                        "未满足约束" + ",差值为" + punish;
                                result.setMessage(message);
                                return result;
                            }
                        } else if (!constrain.getCompareField().equals(FieldEnum.QOUT_DIFF)) {
                            boolean checkResult = CommonFitness.checkField(logFieldDto, checkData, constrain,
                                    topoObjectDto.getFpIndex().getFloodCharacter());
                            if (!checkResult) {
                                result.setFlg(false);
                                String message = constrain.getStnm() + "(" + constrain.getStcd() + ")" +
                                        constrain.getCompareField().getName() + constrain.getLogic().getName() +
                                        "未满足约束";
                                result.setMessage(message);
                                return result;
                            }
                        }
                    }
                }
            }
        }
        result.setFlg(true);
        return result;
    }
}

